# 环节4：构建Milvus向量数据库（适配pymilvus 2.6.0，移除get_client_version）
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from pymilvus import (
    connections,
    FieldSchema, CollectionSchema, DataType,
    Collection,
    utility,
    MilvusException
)
import os
import uuid
import time

# Milvus基础配置
MILVUS_HOST = "localhost"
MILVUS_PORT = "19530"
COLLECTION_NAME = "insurance_manual"
CONNECT_TIMEOUT = 30
RETRY_COUNT = 3


def initialize_tongyi_embeddings(api_key, model_name="text-embedding-v4"):
    """初始化通义千问向量化工具，自动检测向量维度"""
    try:
        embeddings = DashScopeEmbeddings(
            model=model_name,
            dashscope_api_key=api_key
        )
        # 测试向量维度
        test_vector = embeddings.embed_query("测试向量维度")
        actual_dim = len(test_vector)
        print(f"✅ 通义千问模型初始化成功：{model_name}（向量维度：{actual_dim}）")
        return embeddings, actual_dim
    except Exception as e:
        print(f"❌ 通义千问初始化失败：{str(e)}")
        raise


def load_split_vectorize_pdf(pdf_path, api_key, model_name="text-embedding-v4"):
    """PDF加载→分割→向量化（自动适配向量维度）"""
    # 1. 加载PDF
    try:
        loader = PyPDFLoader(pdf_path)
        pages = loader.load()
        print(f"📄 成功加载PDF：{pdf_path}（{len(pages)}页）")
    except FileNotFoundError:
        raise FileNotFoundError(f"❌ 找不到PDF文件：{pdf_path}")

    # 2. 分割文档
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=100,
        length_function=len
    )
    split_docs = splitter.split_documents(pages)
    print(f"✂️  文档分割完成：{len(split_docs)}个片段")

    # 3. 向量化（自动检测维度）
    embeddings, vector_dim = initialize_tongyi_embeddings(api_key, model_name)
    texts = [doc.page_content for doc in split_docs]
    total = len(texts)
    batch_size = 10
    all_vectors = []

    print(f"📡 开始调用通义千问API（分{total // batch_size + 1}批处理，向量维度：{vector_dim}）...")
    for i in range(0, total, batch_size):
        batch_texts = texts[i:i + batch_size]
        batch_num = i // batch_size + 1
        print(f"   处理批次 {batch_num}：{len(batch_texts)}个片段（{i + 1}-{min(i + batch_size, total)}/{total}）")

        try:
            batch_vectors = embeddings.embed_documents(batch_texts)
            # 校验批次向量维度
            for vec in batch_vectors:
                if len(vec) != vector_dim:
                    raise ValueError(f"❌ 批次{batch_num}向量维度错误：应为{vector_dim}维，实际{len(vec)}维")
            all_vectors.extend(batch_vectors)
        except Exception as e:
            raise Exception(f"❌ 批次{batch_num}处理失败：{str(e)}")

    # 校验总数量
    if len(all_vectors) != len(split_docs):
        raise ValueError(f"❌ 向量与片段数量不匹配：向量{len(all_vectors)}个，片段{len(split_docs)}个")

    print(f"✅ 向量化完成：{len(all_vectors)}个{vector_dim}维向量")
    return split_docs, all_vectors, vector_dim


def create_milvus_collection(collection_name, vector_dim):
    """创建Milvus集合（适配pymilvus 2.6.0）"""
    # 定义字段
    fields = [
        FieldSchema(name="id", dtype=DataType.VARCHAR, max_length=64, is_primary=True),
        FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=vector_dim),
        FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=2000),
        FieldSchema(name="page", dtype=DataType.INT64)
    ]

    schema = CollectionSchema(
        fields=fields,
        description=f"平安保险用户手册向量集合（维度：{vector_dim}）"
    )

    # 删除旧集合
    if utility.has_collection(collection_name):
        try:
            utility.drop_collection(collection_name)
            print(f"⚠️  旧集合{collection_name}已删除，创建新集合")
        except MilvusException as e:
            raise Exception(f"❌ 删除旧集合失败：{e.message}")

    # 创建新集合
    try:
        collection = Collection(
            name=collection_name,
            schema=schema,
            using='default'
        )
        print(f"✅ 集合创建完成：{collection_name}（向量维度：{vector_dim}）")
    except MilvusException as e:
        raise Exception(f"❌ 集合创建失败：{e.message}")

    # 创建索引（适配无返回值API）
    index_params = {
        "index_type": "IVF_FLAT",
        "metric_type": "L2",
        "params": {"nlist": 128}
    }
    try:
        collection.create_index(
            field_name="vector",
            index_params=index_params
        )
        print(f"✅ 索引创建成功（IVF_FLAT/L2，nlist=128）")
    except MilvusException as e:
        raise Exception(f"❌ 索引创建失败：{e.message}")

    return collection


# --------------------------
# 核心修改：移除get_client_version()调用
# --------------------------
def connect_milvus_with_retry():
    """带重试的Milvus连接（适配pymilvus 2.6.0，无get_client_version）"""
    for retry in range(RETRY_COUNT):
        try:
            # 断开旧连接
            if connections.has_connection("default"):
                connections.disconnect("default")
                print(f"ℹ️  断开旧连接，重试{retry + 1}/{RETRY_COUNT}")

            # 连接Milvus（无返回值，通过后续获取服务版本验证）
            connections.connect(
                alias="default",
                host=MILVUS_HOST,
                port=MILVUS_PORT,
                timeout=CONNECT_TIMEOUT
            )

            # 验证连接：仅获取服务版本（客户端版本不再获取，适配2.6.0）
            try:
                server_version = utility.get_server_version()
                print(f"✅ Milvus连接成功！服务版本：{server_version}（客户端版本：pymilvus 2.6.0）")
                return True
            except Exception as e:
                raise Exception(f"连接已建立，但获取服务版本失败：{str(e)}")

        except Exception as e:
            print(f"⚠️  连接失败（重试{retry + 1}/{RETRY_COUNT}）：{str(e)}")
            if retry < RETRY_COUNT - 1:
                time.sleep(5)  # 重试前等待5秒
            else:
                raise Exception(f"❌ Milvus连接失败（已重试{RETRY_COUNT}次）：{str(e)}")


def build_milvus_vector_db(split_docs, vectors, vector_dim, collection_name=COLLECTION_NAME):
    """构建Milvus向量数据库（适配pymilvus 2.6.0）"""
    # 校验向量维度
    for i, vec in enumerate(vectors):
        if len(vec) != vector_dim:
            raise ValueError(f"❌ 第{i + 1}个向量维度错误：应为{vector_dim}维，实际{len(vec)}维")

    # 连接Milvus
    connect_milvus_with_retry()

    # 创建集合
    collection = create_milvus_collection(collection_name, vector_dim)

    # 准备插入数据
    data = []
    for i in range(len(split_docs)):
        doc = split_docs[i]
        data.append({
            "id": str(uuid.uuid4()),
            "vector": vectors[i],
            "text": doc.page_content[:1999],  # 截断避免超长
            "page": doc.metadata.get("page", 0)
        })

    # 批量插入（适配2.6.0的MutationResult格式）
    try:
        mr = collection.insert(data)
        # 验证插入结果：通过primary_keys数量判断
        if not hasattr(mr, "primary_keys") or len(mr.primary_keys) == 0:
            raise Exception(f"❌ 数据插入失败：未返回有效主键")

        inserted_count = len(mr.primary_keys)
        print(f"✅ 数据插入成功：{inserted_count}条记录（向量维度：{vector_dim}）")
    except MilvusException as e:
        raise Exception(f"❌ 插入数据时出错：{e.message}")
    except Exception as e:
        raise Exception(f"❌ 插入数据时未知错误：{str(e)}")

    # 加载集合到内存
    try:
        collection.load()
        # 验证加载结果：通过num_entities确认
        entity_count = collection.num_entities
        print(f"✅ 集合加载完成，总记录数：{entity_count}")
    except MilvusException as e:
        raise Exception(f"❌ 集合加载失败：{e.message}")

    return collection, vector_dim


def load_milvus_collection(collection_name=COLLECTION_NAME):
    """加载Milvus集合（适配pymilvus 2.6.0）"""
    # 连接Milvus
    connect_milvus_with_retry()

    # 检查集合是否存在
    if not utility.has_collection(collection_name):
        raise FileNotFoundError(f"❌ 集合{collection_name}不存在，请先构建数据库")

    # 加载集合
    try:
        collection = Collection(collection_name)
        collection.load()
        # 获取向量维度
        vector_field = next(f for f in collection.schema.fields if f.name == "vector")
        vector_dim = vector_field.params["dim"]
        # 验证记录数
        entity_count = collection.num_entities
        print(f"✅ 加载集合完成：{collection_name}（记录数：{entity_count}，维度：{vector_dim}）")
        return collection, vector_dim
    except MilvusException as e:
        raise Exception(f"❌ 集合加载失败：{e.message}")


def search_milvus(collection, query_vector, vector_dim, top_k=2):
    """Milvus相似性检索（适配pymilvus 2.6.0）"""
    # 校验查询向量维度
    if len(query_vector) != vector_dim:
        raise ValueError(f"❌ 查询向量维度错误：应为{vector_dim}维，实际{len(query_vector)}维")

    # 检索参数
    search_params = {
        "metric_type": "L2",
        "params": {"nprobe": 10}
    }

    # 执行检索
    try:
        results = collection.search(
            data=[query_vector],
            anns_field="vector",
            param=search_params,
            limit=top_k,
            output_fields=["text", "page"]
        )
    except MilvusException as e:
        raise Exception(f"❌ 检索失败：{e.message}")

    # 解析结果
    relevant_docs = []
    try:
        for result in results:
            # 验证结果状态（部分版本保留status属性）
            if hasattr(result, "status") and not result.status.OK():
                raise Exception(f"❌ 检索结果无效：{result.status.message}")
            for hit in result:
                relevant_docs.append({
                    "id": hit.id,
                    "distance": hit.distance,
                    "text": hit.entity.get("text", "") if hasattr(hit, "entity") else "",
                    "page": hit.entity.get("page", 0) + 1 if hasattr(hit, "entity") else 0
                })
    except Exception as e:
        raise Exception(f"❌ 结果解析失败：{str(e)}")

    return relevant_docs


if __name__ == "__main__":
    # 配置参数（请替换为你的信息）
    PDF_FILE_PATH = "平安保险用户手册.pdf"
    TONGYI_API_KEY = "sk-6434b554122c4b049ceb805d703f695b"
    TONGYI_MODEL = "text-embedding-v4"  # 实际返回1024维，按API返回为准

    try:
        # 1. PDF加载→分割→向量化
        print("=" * 50)
        print("1. 开始：PDF加载→分割→向量化")
        print("=" * 50)
        split_docs, doc_vectors, vector_dim = load_split_vectorize_pdf(
            pdf_path=PDF_FILE_PATH,
            api_key=TONGYI_API_KEY,
            model_name=TONGYI_MODEL
        )

        # 2. 构建Milvus数据库
        print("\n" + "=" * 50)
        print("2. 开始：构建Milvus向量数据库")
        print("=" * 50)
        milvus_collection, vector_dim = build_milvus_vector_db(
            split_docs=split_docs,
            vectors=doc_vectors,
            vector_dim=vector_dim
        )

        # 3. 测试检索
        print("\n" + "=" * 50)
        print("3. 测试：检索与「平安福投保条件」相关的片段")
        print("=" * 50)
        loaded_collection, vector_dim = load_milvus_collection()

        # 生成查询向量
        query = "平安福2024版的投保条件是什么？"
        embeddings, _ = initialize_tongyi_embeddings(TONGYI_API_KEY, TONGYI_MODEL)
        query_embedding = embeddings.embed_query(query)

        # 执行检索
        results = search_milvus(
            collection=loaded_collection,
            query_vector=query_embedding,
            vector_dim=vector_dim,
            top_k=2
        )

        # 打印结果
        print(f"🔍 检索关键词：{query}")
        if len(results) == 0:
            print("⚠️  未找到相关片段")
        else:
            for i, doc in enumerate(results, 1):
                print(f"\n【相关片段 {i}】")
                print(f"📄 来源页码：第 {doc['page']} 页")
                text_preview = doc['text'][:150] + "..." if len(doc['text']) > 150 else doc['text']
                print(f"📝 片段内容：{text_preview}")
                print(f"📊 相似度距离：{doc['distance']:.4f}（数值越小越相关）")

        # 释放集合内存
        loaded_collection.release()
        print(f"\n✅ 集合已释放内存，程序运行完成！")

    except Exception as e:
        print(f"\n❌ 程序运行失败：{str(e)}")